package com.raymond.algorithm.currentlimit;

import com.google.common.util.concurrent.AtomicDouble;
import com.google.common.util.concurrent.RateLimiter;

/**
 * 令牌桶限流
 * 一.思想
 * 我们有一个固定容量的桶，桶里存放着令牌（token）。
 * 桶一开始是空的，token以 一个固定的速率r往桶里填充，
 * 直到达到桶的容量，多余的令牌将会被丢弃。每当一个请求过来时，
 * 就会尝试从桶里移除一个令牌，如果没有令牌的话，请求无法通过
 *
 * 二.并发数
 * 桶里没有令牌的情况：令牌存放的固定速率
 * 桶里有令牌的情况:桶中的令牌 + 令牌存放的固定速率
 *
 * 三.实现
 * 定义一个桶的大小
 * 定义令牌存放的速率
 * 当来请求的时候,桶中有令牌,直接拿走一个令牌
 * 桶里没有令牌的情况,拿不到令牌,等待
 *
 * 桶中的令牌计算公式:桶中的令牌 = (当前时间 - 上次存放令牌时间) * 令牌存放的速率
 *
 *
 * @author :  raymond
 * @version :  V1.0
 * @date :  2021-04-25 15:52
 */
public class TokenBucket {

    private AtomicDouble token = new AtomicDouble(0);
    /**
     * 每秒的流速(也就是桶的大小)
     */
    private int permitsPerSecond;
    /**
     * 当前时间
     */
    private double time = System.currentTimeMillis();
    /**
     * 单个令牌产生的时间(毫秒)
     */
    private double velocity;

    /**
     * 1秒等于1000毫秒
     */
    public static final int SECONDS = 1000;
    /**
     * 令牌桶
     * @param permitsPerSecond 桶的最大容量,也就是每秒最多参数令牌数
     */
    public TokenBucket(int permitsPerSecond) {
        this.permitsPerSecond = permitsPerSecond;
        this.velocity = 1000D / permitsPerSecond;
    }

    /**
     * 从新设置每秒请求数
     * @param permitsPerSecond 每秒请求数
     */
    public void setPermitsPerSecond(int permitsPerSecond) {
        this.permitsPerSecond = permitsPerSecond;
        this.velocity = 1000D / permitsPerSecond;
    }

    /**
     * 获取一个令牌 不够休眠到能获取为止
     * @throws InterruptedException 休眠异常
     *
     */
    public void acquire() throws InterruptedException {
        acquire(1);
    }


    /**
     * 获取多个令牌，不够到获取为止
     * @param permits 令牌数
     * @throws InterruptedException 休眠异常
     */
    private void acquire(int permits) throws InterruptedException {
        int residueToken = getTokenCount(permits);
        if (residueToken >= permits) {
            return;
        }
        int needPermits = permits - residueToken;
        sleep(needPermits);
    }

    /**
     * 令牌需要休眠时间
     * 休眠后会重新获取令牌
     * @param needPermits 剩余令牌数
     * @throws InterruptedException 休眠异常
     */
    private void sleep(int needPermits) throws InterruptedException {
        Thread.sleep(getSleep(needPermits));
        acquire(needPermits);
    }

    /**
     * 通过令牌数获取休眠时间
     * @param needPermits 令牌数
     * @return 休眠时间
     */
    private long getSleep(int needPermits) {
        double sleep = needPermits * velocity;
        sleep = Math.ceil(sleep);
        if (sleep > SECONDS) {
            sleep = SECONDS;
        }
        return (long) sleep;
    }

    /**
     * 获取需要的令牌数
     * 令牌不够会重新计算令牌数
     * 最后不够获取剩余令牌数据
     * @param count 需要的令牌数
     * @return 令牌数
     */
    public int getTokenCount(int count) {
        if (getBucketToken(count) > 0) {
            return count;
        }
        afreshCalculateToken();
        return getResidueToken(count);
    }

    /**
     * 计算当前时间和上次获取时间来更新令牌桶中的令牌数
     */
    private synchronized void afreshCalculateToken() {
        double now = System.currentTimeMillis();
        double generationGap = now - time;
        double currentCount = generationGap / velocity;
        if (currentCount >= 1) {
            if (currentCount > permitsPerSecond) {
                token.lazySet(permitsPerSecond);
            } else {
                token.addAndGet(currentCount);
            }
            time = now;
        }
    }

    /**
     * 获取桶中剩余令牌数
     * 令牌不够获取剩余令牌数据
     * @param count 需要获取的令牌数
     * @return 令牌数
     */
    private int getResidueToken(int count) {
        if (getBucketToken(count) > 0) {
            return count;
        }
        return getResidueToken();
    }

    /**
     * 获取令牌桶中的剩余数量
     * @return 令牌桶中的剩余数量
     */
    private int getResidueToken() {
        double current = token.get();
        do {
            int currentToken = (int)current;
            if (token.compareAndSet(current, current - currentToken)) {
                return currentToken;
            }
            current = token.get();
        } while (current >= 1);
        return 0;
    }

    /**
     * 获取想要的令牌数
     * @param count 令牌数
     * @return 令牌数，不够为0
     */
    private int getBucketToken(int count) {
        double current = token.get();
        while (current >= count) {
            if (token.compareAndSet(current, current - count)) {
                return count;
            }
            current = token.get();
        }
        return 0;
    }


    /**
     * 初版获取令牌
     * @throws InterruptedException 线程休眠异常
     */
    private void firstEdition() throws InterruptedException {
        if (token.get() >= 1) {
            token.addAndGet(-1);
            return;
        }
        //获取当前时间
        double now = System.currentTimeMillis();
        //获取间隔时间
        double generationGap = now - time;
        //间隔时间除以每个令牌存放时间得到时间段的令牌数
        double count = generationGap / velocity;
        if (count > 1) {
            if (count > permitsPerSecond) {
                count = permitsPerSecond;
            }
            token.addAndGet(count - 1);

        } else {
            //不足一个则计算满足一个令牌需要多少时间
            //用一个令牌需要的时间减去当前已间隔的时间
            //得到一个令牌剩余需要休眠多久
            double sleep = velocity - generationGap;
            if (sleep < 1) {
                sleep = Math.ceil(sleep);
                //休眠时间除以令牌生成时间得到令牌个数
                //令牌个数减去当前需要获取个令牌数
                token.addAndGet(sleep / velocity - 1);
            }
            //更新时间
            now = now + sleep;
            //休眠
            Thread.sleep((long) sleep);
        }
        time = now;
    }

    public static void main(String[] args) throws InterruptedException {
//        gTokenBucket();
        myTokenBucket();

    }

    private static void myTokenBucket() throws InterruptedException {

        TokenBucket rateLimiter = new TokenBucket(1);

        long start = System.currentTimeMillis();
        for (int i = 0; i < 3; i++) {
            rateLimiter.acquire(3);
            if (i == 100) {
//               Thread.sleep(500);
            }

//            int count;
//            if ((count = rateLimiter.getTokenCount(100)) > 0) {
//                System.out.println(count);
//                continue;
//            }
//            break;
//            if (i == 200) {
//                rateLimiter.setPermitsPerSecond(300);
//            }
//            System.out.println(i);
//            System.out.println(System.currentTimeMillis() - start);
        }
        System.out.println(System.currentTimeMillis() - start);
    }

    private static void gTokenBucket() throws InterruptedException {
        AtomicDouble atomicDouble = new AtomicDouble();
        atomicDouble.addAndGet(-1);

        RateLimiter rateLimiter = RateLimiter.create(20000);

        long start = System.currentTimeMillis();
        for (int i = 0; i < 40001; i++) {
            rateLimiter.acquire();
//            if (i == 0) {
//                Thread.sleep(200);
//            }
//            System.out.println(i);
//            System.out.println(System.currentTimeMillis() - start);
//            rateLimiter.acquire()
            System.out.println(rateLimiter.getRate());
        }
        System.out.println(System.currentTimeMillis() - start);
    }
}
